package netfilter

import (
	"fmt"
	"net"

	"gitlab.xunlei.cn/xacc/fastlogger"
	"gvisor.dev/gvisor/pkg/tcpip"
	"gvisor.dev/gvisor/pkg/tcpip/header"
)

var OpenFilterLog = false
var FilterLog = func(fmtStr string, args ...interface{}) {
	if !OpenFilterLog {
		return
	}
	fastlogger.Tracef("[Filter] "+fmtStr, args...)
}

// 使用tcpip包创建特定的私网地址子网来检查
func isInPrivateSubnet(addr tcpip.Address) bool {
	if addr.Len() == 4 {
		// IPv4私网子网
		privateSubnets := []tcpip.Subnet{
			// 10.0.0.0/8
			mustCreateSubnet(tcpip.AddrFromSlice([]byte{10, 0, 0, 0}), 8),
			// 172.16.0.0/12
			mustCreateSubnet(tcpip.AddrFromSlice([]byte{172, 16, 0, 0}), 12),
			// 192.168.0.0/16
			mustCreateSubnet(tcpip.AddrFromSlice([]byte{192, 168, 0, 0}), 16),
			// 127.0.0.0/8
			mustCreateSubnet(tcpip.AddrFromSlice([]byte{127, 0, 0, 0}), 8),
			// 169.254.0.0/16
			mustCreateSubnet(tcpip.AddrFromSlice([]byte{169, 254, 0, 0}), 16),
		}

		for _, subnet := range privateSubnets {
			if subnet.Contains(addr) {
				return true
			}
		}
	} else if addr.Len() == 16 {
		// IPv6私网子网
		// fe80::/10 (链路本地地址)
		linkLocalSubnet := mustCreateSubnet(
			tcpip.AddrFromSlice([]byte{0xfe, 0x80, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}),
			10,
		)
		if linkLocalSubnet.Contains(addr) {
			return true
		}
		fastlogger.Infof("addr: %s", addr.String())
		// fc00::/7 (唯一本地地址)
		ulaSubnet := mustCreateSubnet(
			tcpip.AddrFromSlice([]byte{0xfc, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}),
			7,
		)
		if ulaSubnet.Contains(addr) {
			return true
		}

		// ::1/128 (本地回环)
		loopbackAddr := tcpip.AddrFromSlice([]byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1})
		if addr == loopbackAddr {
			return true
		}
	}
	return false
}

// 辅助函数：创建子网（带错误处理）
func mustCreateSubnet(addr tcpip.Address, prefixLen int) tcpip.Subnet {
	// 创建掩码
	addrLen := addr.Len()
	mask := make([]byte, addrLen)

	// 设置前缀位为1
	for i := 0; i < prefixLen/8; i++ {
		mask[i] = 0xff
	}
	if prefixLen%8 != 0 {
		mask[prefixLen/8] = byte(0xff << (8 - prefixLen%8))
	}

	// 应用掩码到地址
	maskedAddr := make([]byte, addrLen)
	for i := 0; i < addrLen; i++ {
		maskedAddr[i] = addr.AsSlice()[i] & mask[i]
	}

	subnet, err := tcpip.NewSubnet(
		tcpip.AddrFromSlice(maskedAddr),
		tcpip.MaskFromBytes(mask),
	)
	if err != nil {
		panic(fmt.Sprintf("Failed to create subnet: %v", err))
	}

	return subnet
}

// 更现代的实现方式，使用tcpip包的AddressWithPrefix
func filterPrivateAddressesModern(ipPkg []byte, version uint8) Verdict {
	switch version {
	case 4: // IPv4
		dstIP := tcpip.AddrFromSlice(ipPkg[16:20])
		if !isInPrivateSubnet(dstIP) {
			return NF_STOLEN
		}

	case 6: // IPv6
		dstIP := tcpip.AddrFromSlice(ipPkg[24:40])

		if !isInPrivateSubnet(dstIP) {
			return NF_STOLEN
		}
	default:
		return NF_ACCEPT
	}
	return NF_ACCEPT
}

type FilterFn func(ipPkg []byte, version int) Verdict

type IpFilter struct {
	FilterFns []FilterFn
}

func (f *IpFilter) Filter(ipPkg []byte, version int) Verdict {
	for _, fn := range f.FilterFns {
		if fn(ipPkg, version) == NF_DROP {
			return NF_DROP // 程序处理
		}
	}
	return NF_ACCEPT // 协议栈处理
}

func (f *IpFilter) AddFilterFn(fn FilterFn) {
	// copy old
	old := f.FilterFns
	// add new
	new := make([]FilterFn, 0, len(old)+1)
	new = append(new, old...)
	new = append(new, fn)
	// swap
	f.FilterFns = new
}

func NewIpFilter() *IpFilter {
	return &IpFilter{
		FilterFns: make([]FilterFn, 0),
	}
}

var PublicDstIp = func(ipPkg []byte, version uint8) Verdict {
	// 默认允许所有流量
	return filterPrivateAddressesModern(ipPkg, version)
}

var IgnoreAllIp = func(ipPkg []byte, version uint8) Verdict {
	return NF_ACCEPT
}

var DefaultIpFilter = NewIpFilter()

type IPv4PassListMap struct {
	ipMap map[uint32]bool
}

var PassTargetIpv4 = func(ips []string) NetFilterRecvIpPacket {
	passMap := NewIPv4PassListMap(ips)
	return func(ipPkg []byte, version uint8) Verdict {
		if version != 4 {
			return NF_ACCEPT
		}
		srcIP := header.IPv4(ipPkg).SourceAddress()
		dstIP := header.IPv4(ipPkg).DestinationAddress()
		if passMap.ContainsAddress(dstIP) {
			FilterLog("NF_STOLEN PassTargetIpv4: %s <-> %s", srcIP.String(), dstIP.String())
			return NF_STOLEN
		}
		return NF_ACCEPT
	}
}

var WatchAllIp = func(ipPkg []byte, version uint8) Verdict {
	var src, dst tcpip.Address
	var protocol uint8
	var l4Pkg []byte
	switch version {
	case 4:
		ipv4 := header.IPv4(ipPkg)
		src = ipv4.SourceAddress()
		dst = ipv4.DestinationAddress()
		protocol = uint8(ipv4.TransportProtocol())
		l4Pkg = ipv4.Payload()
	case 6:
		ipv6 := header.IPv6(ipPkg)
		src = ipv6.SourceAddress()
		dst = ipv6.DestinationAddress()
		protocol = uint8(ipv6.NextHeader())
		l4Pkg = ipv6.Payload()
	default:
		return NF_ACCEPT
	}
	switch protocol {
	case uint8(header.TCPProtocolNumber):
		tcp := header.TCP(l4Pkg)
		FilterLog("TCP %s:%d -> %s:%d", src.String(), tcp.SourcePort(), dst.String(), tcp.DestinationPort())
	case uint8(header.UDPProtocolNumber):
		udp := header.UDP(l4Pkg)
		FilterLog("UDP %s:%d -> %s:%d", src.String(), udp.SourcePort(), dst.String(), udp.DestinationPort())
	case uint8(header.ICMPv4ProtocolNumber):
		icmp := header.ICMPv4(l4Pkg)
		FilterLog("ICMP %s -> %s, type=%d code=%d", src.String(), dst.String(), icmp.Type(), icmp.Code())
	case uint8(header.ICMPv6ProtocolNumber):
		icmpv6 := header.ICMPv6(l4Pkg)
		FilterLog("ICMPv6 %s -> %s, type=%d code=%d", src.String(), dst.String(), icmpv6.Type(), icmpv6.Code())
	default:
		FilterLog("Other Protocol %s -> %s, protocol=%d", src.String(), dst.String(), protocol)
	}
	return NF_ACCEPT
}

// IP字符串转uint32
func IpStringToUint32(ipStr string) (uint32, error) {
	ip := net.ParseIP(ipStr)
	if ip == nil {
		return 0, fmt.Errorf("invalid IP address: %s", ipStr)
	}

	ip = ip.To4()
	if ip == nil {
		return 0, fmt.Errorf("not an IPv4 address: %s", ipStr)
	}

	return uint32(ip[0])<<24 | uint32(ip[1])<<16 | uint32(ip[2])<<8 | uint32(ip[3]), nil
}

func NewIPv4PassListMap(ipStrings []string) *IPv4PassListMap {
	passlist := &IPv4PassListMap{
		ipMap: make(map[uint32]bool),
	}

	for _, ipStr := range ipStrings {
		if ipUint32, err := IpStringToUint32(ipStr); err == nil {
			passlist.ipMap[ipUint32] = true
		}
	}

	return passlist
}

// O(1) 查找
func (p *IPv4PassListMap) Contains(ip uint32) bool {
	return p.ipMap[ip]
}

func (p *IPv4PassListMap) ContainsAddress(addr tcpip.Address) bool {
	if addr.Len() != 4 {
		return false
	}

	ip := addr.As4()
	ipUint32 := uint32(ip[0])<<24 | uint32(ip[1])<<16 | uint32(ip[2])<<8 | uint32(ip[3])
	return p.Contains(ipUint32)
}
